"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2024/7/15 19:37
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    compare_ecovacs_sample_alg_effect.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    compare_ecovacs_sample_alg_effect_VER0.1

------------------------------------------------------------------
# @Description:
    TODO: 注意，该脚本寻找对应的rgb数据，并将其重新整理并生成深度效果，然后计算新的点云
------------------------------------------------------------------
"""
import os
import cv2
import numpy as np
from utils import Utils
from utils import Colorizer
from tqdm import tqdm
from natsort import natsorted
from align import RMSLAlign
from parser_cam_param import ParserCamParam


def findMatchedImagePath2Txt(test_dir: str, find_dir: str, txt_path: str, key='_color'):
    if not os.path.exists(find_dir) or not os.path.exists(test_dir):
        print("[Error] Path not exists!")
        return

    cur_rgb_list = Utils().findFilesWithPathlibRecursive(test_dir, pattern1=key, pattern2='.png')
    all_rgb_list = Utils().findFilesWithPathlibRecursive(find_dir, pattern1=key, pattern2='.png')

    dict_cur = {}
    dict_all = {}

    # Generate hash values for current directory images
    for index, rgb_path in enumerate(
            tqdm(cur_rgb_list, desc='Generating hash values for test directory...', colour='cyan')):
        dict_cur[rgb_path] = Utils().get_file_hash(rgb_path)

    # Generate hash values for all directory images
    for index, rgb_path in enumerate(
            tqdm(all_rgb_list, desc='Generating hash values for find directory...', colour='cyan')):
        dict_all[rgb_path] = Utils().get_file_hash(rgb_path)

    # Find matching paths
    matched_paths = {}
    for cur_path, cur_hash in dict_cur.items():
        matched_list = []
        for all_path, all_hash in dict_all.items():
            if cur_hash == all_hash and cur_path != all_path:
                matched_list.append(all_path)
        if matched_list:
            matched_paths[cur_path] = matched_list

    # Write matches to a txt file
    with open(txt_path, 'w') as f:
        for cur_path, matched_list in matched_paths.items():
            f.write(f'{cur_path} -> {matched_list}\n')
        print(f"Matching results saved to {txt_path}")


def generatePointCloudUseTxt(testDir: str, txt_path: str, key0='rk_result_0710/', key1='_disp', key2='_depth', key4='_color', key5='_speckle-left'):
    all_disp_path_list = Utils().findFilesWithPathlibRecursive(testDir, pattern1=key1, pattern2='.png')
    all_disp_path_list = [img_path for img_path in all_disp_path_list if '_align' not in img_path]
    all_disp_path_list = natsorted(all_disp_path_list)

    # Generate all test array to visual studio
    all_res_path_list = []
    for idx, disp_path in enumerate(tqdm(all_disp_path_list, desc=' Generating test dataset!', colour='cyan')):
        rgb_img_path = disp_path.replace(key0, '').replace(key1, key4)
        depth_img_path = disp_path.replace(key1, key2)
        ir_img_path = rgb_img_path.replace(key4, key5)

        all_res_path_list.append(rgb_img_path)
        all_res_path_list.append(disp_path)
        all_res_path_list.append(depth_img_path)
        all_res_path_list.append(ir_img_path)

    with open(txt_path, 'w') as f:
        for img_path in all_res_path_list:
            f.write(f'{img_path}\n')
        print(f"Done!")


if __name__ == '__main__':
    test_match_dir = '/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/r1920x1080_d1280x800/compare_data'
    find_match_dir = '/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/r1920x1080_d1280x800'
    save_txt_dir = '/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/r1920x1080_d1280x800/compare_data/match_res.txt'

    # findMatchedImagePath2Txt(test_match_dir, find_match_dir, save_txt_dir, key='_color')

    test_dir_path = "/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/r1920x1080_d1280x800/rk_result_0710"
    save_txt_path = "/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/r1920x1080_d1280x800/rk_result_0710/vs_use.txt"
    generatePointCloudUseTxt(test_dir_path, save_txt_path, key0='rk_result_0710', key1='_disp', key2='_depth', key4='_color', key5='_speckle-left')
